
import matplotlib.pyplot as plt

class GUIHelper:

    @staticmethod
    def createVolSurf3DFigure(volSurface):

        fig = plt.figure()
        ax = fig.gca(projection='3d')
        X, Y, Z = volSurface.getDiscretizedXYZTuples()
        ax.plot_surface(X, Y, Z, rstride=8, cstride=8, alpha=0.3)

        #ax.contour(X, Y, Z, zdir='z', offset=0.5)
        ax.contour(X, Y, Z, zdir='x', offset=0.5)
        ax.contour(X, Y, Z, zdir='y', offset=0.5)

#        ax.contour(X, Y, Z, zdir='z', offset=-100)
#        ax.contour(X, Y, Z, zdir='x', offset=-40)
#        ax.contour(X, Y, Z, zdir='y', offset=40)

        ax.set_xlabel('K/S')
#        ax.set_xlim(-40, 40)
        ax.set_ylabel('Term')
#        ax.set_ylim(-40, 40)
        ax.set_zlabel('Volatility')
#        ax.set_zlim(-100, 100)

        ax.set_title('Volatility Surface')

        return fig



    @staticmethod
    def createVolCurveFit2DFigure(volSurface, expDate, optionMktData):
        fig = plt.figure()

        #plot data
        ax = fig.add_subplot(1,1,1)
        X, Y = volSurface.getDiscretizedXYTuples(expDate)
        ax.plot(X,Y,'--')

        optionPriceChain = optionMktData.optionPriceChainMap[expDate]

#        X, Y = optionPriceChain.getXYTuplesForCalls()
#        ax.plot(X,Y,'bo')
#
#        X, Y = optionPriceChain.getXYTuplesForPuts()
#        ax.plot(X,Y,'bo')

        #ax.set_xlim(5,0)

        ax.set_xlabel('K/S')
        ax.set_ylabel('Volatility')
        ax.set_title('Expiration Curve')

        return fig